'''
Utility variables and functions for program action needs
'''

import os

from PyQt4 import QtGui

from util_background import *
from class_items import *

AA_LIST_WITH_GAP = 'ACDEFGHIKLMNPQRSTVWY-'


def readSequenceFile(fin):
    """
    Read sequences from a file
    Acceptable file formats: .fasta and .txt
    """
    if fin is None:
        data = None
        
    elif not fin.endswith('.fasta') and not fin.endswith('.txt'):
        messageBox = QtGui.QMessageBox()
        messageBox.setWindowTitle('Error')
        messageBox.setText('Only .fasta or .txt files are supported.')
        messageBox.exec_()
        data = None
        
    elif not os.path.exists(fin):
        data = None
        
    else:
        
        data = {}
        
        with open(fin,  'r') as f:
            
            if fin.endswith('.txt'):
                
                allines = [eachline.strip() for eachline in f.readlines()]
                data['seq']= [eachline for eachline in allines if eachline != '']
                data['num']  = len(data['seq'])
                data['id'] = [''] * data['num']
                data['maxSeqLen'] = max([len(eachSeq) 
                                         for eachSeq in data['seq']])
            
            elif fin.endswith('.fasta'):
                
                allines = f.read()
                id_seq_pairs = [eachpair for eachpair in allines.split('>') 
                                if eachpair != '']
                
                ids = []
                seqs = []
                
                for eachpair in id_seq_pairs:
                    id_seq = eachpair.strip().split('\n')
                    ids.append(id_seq[0])
                    seqs.append(''.join(id_seq[1:]))
                
                data['id'] = ids
                data['seq'] = seqs
                data['num'] = len(seqs)
                data['maxSeqLen'] = max([len(eachSeq) 
                                         for eachSeq in data['seq']])
            
            # validate input sequences
            for seq in data['seq']:
                for aa in seq:
                    if aa not in AA_LIST_WITH_GAP:
                        messageBox = QtGui.QMessageBox()
                        messageBox.setWindowTitle('Error')
                        messageBox.setText('Invalid protein sequences!')
                        messageBox.exec_()
                        return None
    
    return data


def saveViewToFile(scene, fout):
    
    if fout.endswith('.pdf') or fout.endswith('.ps'):
        saveViewToPrintableFile(scene, fout)
        
    else:
        saveViewToImageFile(scene, fout)


def saveViewToImageFile(scene, fout):
    
    image = QtGui.QImage(1600, 1200, 
                         QtGui.QImage.Format_ARGB32_Premultiplied)
    image.fill(0xFFFFFF)
    
    painter = QtGui.QPainter(image)
    painter.setRenderHint(QtGui.QPainter.Antialiasing)
    scene.render(painter)
    painter.end()

    image.save(fout)


def saveViewToPrintableFile(scene, fout):

    printer = QtGui.QPrinter(QtGui.QPrinter.HighResolution)
    
    if fout.endswith('.pdf'):
        printer.setOutputFormat(QtGui.QPrinter.PdfFormat)
    elif fout.endswith('.ps'):
        printer.setOutputFormat(QtGui.QPrinter.PostScriptFormat)
    
    printer.setPaperSize(QtGui.QPrinter.Letter)
    printer.setOrientation(QtGui.QPrinter.Landscape)
    printer.setOutputFileName(fout)
    
    painter = QtGui.QPainter()
    painter.begin(printer)
    painter.setRenderHint(QtGui.QPainter.Antialiasing)
    scene.render(painter)
    painter.end()


def saveSettingsDefault(settings):
    
    saveSettingsToFile(settings, FILE_SETTINGS_DEFAULT)


def saveSettingsRecent(settings):
    
    saveSettingsToFile(settings, FILE_SETTINGS_RECENT)


def saveSettingsToFile(settings, settingsFile):
    
    AAList = 'ACDEFGHIKLMNPQRSTVWY'
    
    fout = open(settingsFile, 'w')
    fout.write('<AA properties>')
    
    for aa in AAList:
        
        fout.write('\n\t<'+aa+'>\n')
        
        for property in settings['aa properties'][aa]:
            valueStr = str(round(settings['aa properties'][aa][property], 2))
            fout.write('\t\t<'+property+'>' + valueStr + '</'+property+'>\n')
        
        fout.write('\t</'+aa+'>\n')
    
    fout.write('</AA properties>\n')
    
    fout.write('\n<property type definition>')
    
    for property in settings['property type definition'].keys():
        
        fout.write('\n\t<'+property+'>\n')
        
        rangeIntervals = list(settings['property type definition'][property].values())
        rangeIntervals.sort()
        
        for rangeInterval in rangeIntervals:
            for eachType in settings['property type definition'][property].keys():
                if settings['property type definition'][property][eachType] == rangeInterval:
                    rangeLB = str(round(rangeInterval[0], 2))
                    rangeUB = str(round(rangeInterval[1], 2))
                    fout.write('\t\t<'+eachType+'>['+rangeLB+', '+\
                                           rangeUB+')</'+eachType+'>\n')
                    break
        
        fout.write('\t</'+property+'>\n')
    
    fout.write('</property type definition>\n')
    
    fout.write('\n<property color mapping>')
    
    for property in settings['property color mapping'].keys():
        
        fout.write('\n\t<'+property+'>\n')
        
        if property == 'identity':
            propertyTypes = sorted(list(settings['property color mapping'][property].keys()))
        else:
            propertyTypes = []
            rangeIntervals = sorted(list(settings['property type definition'][property].values()))
            for rangeInterval in rangeIntervals:
                for eachType in settings['property type definition'][property].keys():
                    if settings['property type definition'][property][eachType] == rangeInterval:
                        propertyTypes.append(eachType)
                        break
        
        for eachType in propertyTypes:            
            colorString = '#' + hex(settings['property color mapping']
                                    [property][eachType].rgb())[4:].upper()
            fout.write('\t\t<'+eachType+'>'+colorString+'</'+eachType+'>\n')
        
        fout.write('\t</'+property+'>\n')
    
    fout.write('</property color mapping>')
    
    fout.close()


def saveAlignmentToFile(seqData, alignmentFile):
    """
    Save alignment sequences to files
    Format supported: .fasta and .txt
    The formats are similar to the input sequence files
    """
    fout = open(alignmentFile, 'w')
    
    if alignmentFile.endswith('.fasta'):
        
        if not seqData['id']:
            assert len(seqData['id']) == len(seqData['seq'])
        
        for i in range(len(seqData['id'])):
            fout.write('>' + seqData['id'][i] + '\n')
            fout.write(seqData['seq'][i] + '\n')
        
    elif alignmentFile.endswith('.txt'):
        
        for seq in seqData['seq']:
            fout.write(seq + '\n')
    
    fout.close()
